
import torch
import time
from ipex_llm.transformers import Qwen2VLForConditionalGeneration
from qwen_vl_utils import process_vision_info
from modelscope import AutoProcessor


model_path = "C:\models\Qwen\Qwen2-VL-2B-Instruct"
# model_path = "C:\models\Qwen\Qwen2.5-VL-7B-Instruct"
model = Qwen2VLForConditionalGeneration.from_pretrained(model_path,
                                                        load_in_4bit=True,
                                                        optimize_model=True,
                                                        trust_remote_code=True,
                                                        modules_to_not_convert=["vision"],
                                                        use_cache=True)

# Use .float() for better output, and use .half() for better speed
model = model.half().to("xpu")

# The following code for generation is adapted from https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct#quickstart

# The default range for the number of visual tokens per image in the model is 4-16384.
# You can set min_pixels and max_pixels according to your needs, such as a token count range of 256-1280,
# to balance speed and memory usage.
min_pixels = 256 * 28 * 28
max_pixels = 1280 * 28 * 28
processor = AutoProcessor.from_pretrained(model_path, min_pixels=min_pixels, max_pixels=max_pixels)

prompt = "描述一下这张图"
image_path = "../examples/demo.jpeg"

with torch.inference_mode():
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image_path,
                },
                {"type": "text", "text": prompt},
            ],
        }
    ]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to('xpu')

    # ipex_llm model needs a warmup, then inference time can be accurate
    generated_ids = model.generate(
        **inputs,
        max_new_tokens=512
    )

    st = time.time()
    generated_ids = model.generate(
        **inputs,
        max_new_tokens=512
    )
    torch.xpu.synchronize()
    end = time.time()
    generated_ids = generated_ids.cpu()
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
    ]

    response = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(f'Inference time: {end - st} s')
    print('-' * 20, 'Input Image', '-' * 20)
    print(image_path)
    print('-' * 20, 'Prompt', '-' * 20)
    print(prompt)
    print('-' * 20, 'Output', '-' * 20)
    print(response)
    time.sleep(5)